import requests, json, hashlib, time, tiktoken, datetime
from json.decoder import JSONDecodeError
from django_redis import get_redis_connection
from django.http import JsonResponse, StreamingHttpResponse
from django.views import View
from django.views.decorators.http import require_POST
from process.models import ChatSession, ChatHistory, Mask
from apikey.models import ApiKey, AllowModel
from user.models import UserProfile
from user.views import ChatSessionToRedis, getChatSessionList
from openai.error import *

tiktoken.get_encoding('cl100k_base')


def get_token_count(message, reply, model):
    encoding = tiktoken.encoding_for_model(model)
    send_token = len(encoding.encode(message['content'])) + 7 \
        if type(message) == dict else\
        sum(len(encoding.encode(item['content'])) + 7 for item in message)
    receive_token = len(encoding.encode(reply)) if reply else 0
    return send_token, receive_token


def AllowSendMessage(user: UserProfile, model):
    Type = user.billing_type
    Member = user.membership_type
    Expired = user.expired
    Now = datetime.datetime.now()
    if Type == 1 and Member == 1 and user.count.get(model, 0) > 0:
        return True
    elif Type == 1 and Member == 2 and user.count.get(model, 0) > 0 and Expired > Now:
        return True
    elif Type == 1 and Member == 3 and user.count.get(model, 0) > 0 and Expired > Now:
        return True
    elif Type == 2 and user.count.get(model, 0) > 0:
        return True
    return False


def update_chat_session_to_redis(userID, chatID, title):
    conn = get_redis_connection('chat_sessions')
    chat_session: bytes = conn.get(userID)
    chat_session_list = json.loads(chat_session.decode())
    chat = next((d for d in chat_session_list if d['chatID'] == chatID), None)
    chat['title'] = title
    conn.set(userID, json.dumps(chat_session_list))


def GPT(userID, chatID, message, model, chat_history: ChatHistory):
    ak = None
    user = chat_history.chatID.userID
    if not AllowSendMessage(user, model):
        yield '暂无权限 或 已没有使用次数'
        return
    history = chat_history.message
    try:
        history.append(message)
        reply = ""
        api_keys = ApiKey.objects.get(model__type=model, state=True)
        ak = api_keys.key
        response = openai.ChatCompletion.create(
            api_key=api_keys.key,
            api_base=api_keys.hostname if api_keys.hostname else "https://api.openai.com/v1",
            model=model,
            messages=history,
            temperature=0.7,
            stream=True
        )
        for line in response:
            if line["finish_reason"] is not None:
                data = ''
            else:
                data = line["choices"][0]["delta"].get("content", "")
                reply += data
            yield data
        if chat_history.title == '新的会话':
            chat_history.title = message['content']
            update_chat_session_to_redis(userID, chatID, message['content'])
        history.append({"role": "assistant", "content": reply})
        chat_history.message = history
        send_token, receive_token = get_token_count(message, reply, model)
        chat_history.send_token += send_token
        chat_history.receive_token += receive_token
        if chat_history.title == '新的会话':
            chat_history.title = message['content']
            update_chat_session_to_redis(user.userID, chat_history.chatID_id, message['content'])
        chat_history.save()
        membership_type = int(user.membership_type)
        billing_type = int(user.billing_type)
        if billing_type == 1 and membership_type == 3 and model == 'gpt-3.5-turbo':
            return user.save()
        value = user.count.get(model, -1)
        if value == -1:
            raise AttributeError
        user.count[model] = value - 1
        user.save()
    except RateLimitError:
        yield '您已达到指定的速率限制，请稍等重试'
    except AuthenticationError:
        ApiKey.objects.filter(key=ak).update(state=False)
        yield 'API 密钥或令牌无效、过期或撤销'
    except APIConnectionError:
        yield '与OpenAI通信时出错，在处理您的请求时出现了问题。这可能是由于临时错误、错误或系统中断造成的。'
    except InvalidRequestError as e:
        error = e.error
        if error.get('code') == 'model_not_found':
            yield f'型号 {model} 不存在，或者您无权访问它'
        elif error.get('code') == 'context_length_exceeded':
            yield '此模型最大上下文长度为4097个Token，请重新创建新的会话'
        else:
            yield error.get('message')
    except ServiceUnavailableError:
        yield "OpenAi服务商暂时无法处理您的请求。这可能是由于计划内或计划外维护、系统升级或服务器故障造成的。或可能在高流量期间暂时无法返回响应。"
    except APIError:
        yield "在处理您的请求时出现了问题，这可能是由于临时错误、错误或系统中断造成的。"
    except ApiKey.DoesNotExist:
        yield f"在 {model} 模型的池中无法找到可用的Key"
    except AttributeError:
        yield "传入参数有误，请重新提交"
    except Exception as e:
        print(type(e), e.__traceback__.tb_lineno)
        print(e)
        yield "请刷新重试……"


@require_POST
def process(request):
    message = request.POST.get("message", None)
    model = request.POST.get("model", None)
    chatID = request.POST.get("chatID", None)
    userID = request.headers.get("userID", None)
    if not all([message, model, userID, chatID]):
        return JsonResponse({"code": 400, "msg": "参数缺失"})
    if model not in [item.type.type for item in AllowModel.objects.all()]:
        return JsonResponse({"code": 404, "msg": "模型不存在"})
    try:
        message = json.loads(message)
        chat_history = ChatHistory.objects.get(chatID__userID=userID, chatID=chatID, chatID__state=True)
        return StreamingHttpResponse(GPT(userID, chatID, message, model, chat_history), content_type='text/event-stream')
    except JSONDecodeError:
        return JsonResponse({"code": 401, "msg": "传入参数类型错误"})
    except ChatHistory.DoesNotExist:
        return JsonResponse({"code": 407, "msg": "聊天记录不存在"})
    except Exception as e:
        print(e.__traceback__.tb_lineno)
        return JsonResponse({"code": 408, "msg": str(e)})


def getChatSession(request):
    userID = request.headers.get("userID", None)
    if not userID:
        return JsonResponse({"code": 404, "msg": "参数缺失"})
    try:
        user = UserProfile.objects.get(userID=userID)
        conn = get_redis_connection('chat_sessions')
        SessionList: bytes = conn.get(user.userID)
        return JsonResponse({"code": 200, "msg": json.loads(SessionList.decode())})
    except UserProfile.DoesNotExist:
        return JsonResponse({"code": 405, "msg": "用户不不存在"})
    except TypeError:
        return JsonResponse({"code": 406, "msg": []})
    except Exception as e:
        return JsonResponse({"code": 500, "msg": str(e)})


def chat_session_id(user_id):
    return hashlib.sha1((user_id + str(int(time.time()))).encode()).hexdigest()[:32]


class ChatSessionView(View):

    def post(self, request):
        userID = request.headers.get("userID", None)
        if not userID:
            return JsonResponse({"code": 404, "msg": "参数缺失"})
        try:
            user = UserProfile.objects.get(userID=userID)
            if ChatSession.objects.filter(userID=user, state=True, chathistory__title='新的会话').exists():
                return JsonResponse({"code": 201, "msg": "已经有一个空的会话了噢"})
            newChatID = chat_session_id(user.userID)
            chat_session = ChatSession.objects.create(userID=user, chatID=newChatID)
            ChatHistory.objects.create(chatID=chat_session)
            ChatSessionToRedis(userID, getChatSessionList(userID))
            return JsonResponse({"code": 200, "msg": newChatID})
        except UserProfile.DoesNotExist:
            return JsonResponse({"code": 405, "msg": "用户不存在"})
        except Exception as e:
            return JsonResponse({"code": 500, "msg": str(e)})

    def delete(self, request):
        userID = request.headers.get("userID", None)
        chatID = json.loads(request.body.decode()).get("chatID", None)
        if not all([userID, chatID]):
            return JsonResponse({"code": 404, "msg": "参数缺失"})
        try:
            session = ChatSession.objects.get(userID=userID, chatID=chatID)
            session.state = False
            session.save()
            ChatSessionToRedis(userID, getChatSessionList(userID))
            return JsonResponse({"code": 200, "msg": "删除成功"})
        except ChatSession.DoesNotExist:
            return JsonResponse({"code": 405, "msg": "会话ID不存在"})
        except Exception as e:
            return JsonResponse({"code": 500, "msg": str(e)})


@require_POST
def getChatHistory(request):
    userID = request.headers.get("userID", None)
    chatID = json.loads(request.body.decode()).get("chatID", None)

    if not all([userID, chatID]):
        return JsonResponse({"code": 404, "msg": "参数缺失"})
    try:
        user = UserProfile.objects.get(userID=userID)
        chat_session = ChatSession.objects.get(userID=user, chatID=chatID, state=True).chathistory_set.first()
        if not chat_session:
            return JsonResponse({"code": 407, "msg": "无法获取该会话的记录"})
        return JsonResponse({"code": 200, "msg": chat_session.message})
    except UserProfile.DoesNotExist:
        return JsonResponse({"code": 405, "msg": "用户不存在"})
    except ChatSession.DoesNotExist:
        return JsonResponse({"code": 406, "msg": "会话ID不存在"})
    except Exception as e:
        return JsonResponse({"code": 500, "msg": str(e)})


def getMask(request):
    try:
        connect = get_redis_connection('mask')
        keys = connect.keys()
        pipe = connect.pipeline()
        for key in keys: pipe.get(key)
        values = pipe.execute()
        List = [{"mask_id": key.decode(), **json.loads(value.decode())} for key, value in zip(keys, values)]
        return JsonResponse({"code": 200, "msg": List})
    except Exception as e:
        return JsonResponse({"code": 500, "msg": str(e)})


def add_session_with_mask(request):
    try:
        userID = request.headers.get('userID', None)
        maskID = json.loads(request.body).get('mask', None)
        if not all([userID, maskID]):
            return JsonResponse({"code": 200, "msg": "缺少参数"})
        connect = get_redis_connection('mask')
        mask = json.loads(connect.get(maskID))
        model = mask.get('model', None)
        message = [{"role": "system", "content": i} for i in mask['desc'].split("|")]
        user = UserProfile.objects.get(userID=userID)
        newChatID = chat_session_id(user.userID)
        chat_session = ChatSession.objects.create(userID=user, chatID=newChatID)
        send_token, receive_token = get_token_count(message, '', model)
        ChatHistory.objects.create(chatID=chat_session, title=f"角色 - {mask['title']}", message=message, send_token=send_token)
        ChatSessionToRedis(userID, getChatSessionList(userID))
        return JsonResponse({"code": 200, "msg": newChatID})
    except UserProfile.DoesNotExist:
        return JsonResponse({"code": 500, "msg": "用户不存在"})
    except Exception as e:
        print(e)
        return JsonResponse({"code": 500, "msg": str(e)})


def MaskToRedis():
    connect = get_redis_connection('mask')
    maskList = Mask.objects.all()
    pipe = connect.pipeline()
    for mask in maskList:
        data = {
            "title": mask.title,
            "count": mask.desc.count("|") + 1,
            "desc": mask.desc,
            "model": mask.get_model_display()
        }
        pipe.set(mask.maskId, json.dumps(data))
    pipe.execute()


from apscheduler.schedulers.background import BackgroundScheduler

scheduler = BackgroundScheduler(timezone='Asia/Shanghai')
scheduler.add_job(MaskToRedis, 'cron', hour=0, name='reset_count_task', id='mask')
scheduler.start()
MaskToRedis()
